import deepxde as dde
import numpy as np
from matplotlib import pyplot as plt

def restore_model(model_path):
    geom = dde.geometry.Rectangle([-5.5, -2], [6.5, 2])
    data = dde.data.PDE(geom, DL_Euler_Equation_2D, [], 0, 0)
    layer_size = [2] + [40] * 6 + [4]
    activation = "tanh"
    initializer = "Glorot uniform"
    net = dde.nn.FNN(layer_size, activation, initializer)
    model = dde.Model(data, net)
    model.compile("adam", lr=0.001)
    model.restore(save_path = model_path)
    return model

"""------------------------------ Dimensionless Euler Equations(Time-dependent PDEs) -----------------------"""
def DL_Euler_Equation_2D(x, y):
    """
    System of PDEs to be minimized: compressible 2D Euler equations.

    """
    gamma = 1.4
    
    u = y[:, 0:1]
    v = y[:, 1:2]
    p = y[:, 2:3]
    rho = y[:, 3:4]

    E = p/(gamma-1) + 0.5*rho*(u*u + v*v)

    A1, B1 = rho*u, rho*v
    A2, B2 = gamma*rho*u*u + p, gamma*rho*u*v
    A3, B3 = gamma*rho*u*v, gamma*rho*v*v + p
    A4, B4 = u*E, v*E

    dA1_x = dde.grad.jacobian(A1, x, i=0, j=0)
    dB1_y = dde.grad.jacobian(B1, x, i=0, j=1)

    dA2_x = dde.grad.jacobian(A2, x, i=0, j=0)
    dB2_y = dde.grad.jacobian(B2, x, i=0, j=1)

    dA3_x = dde.grad.jacobian(A3, x, i=0, j=0)
    dB3_y = dde.grad.jacobian(B3, x, i=0, j=1)

    dA4_x = dde.grad.jacobian(A4, x, i=0, j=0)
    dB4_y = dde.grad.jacobian(B4, x, i=0, j=1)
    
    continuity = dA1_x + dB1_y
    x_momentum = dA2_x + dB2_y
    y_momentum = dA3_x + dB3_y
    energy = dA4_x + dB4_y
    return [continuity, x_momentum, y_momentum, energy]

def plot_naca001065_flow_field(model):
    """------------------------------ Fluid parameters  -----------------------"""
    # 空气绝热指数
    k = 1.4
    # 用于无量纲化的速度、压力和密度，取入口处的声速值、压力值和密度值
    p0 = 73048 #(pa)
    rho0 = 0.9 #(kg/m^3)
    a0 = np.sqrt(k*p0/rho0) #(m/s)

    import matplotlib as mpl
    from matplotlib import rcParams
    config = {"font.size": 16}
    rcParams.update(config)
    plt.rcParams['xtick.direction'] = 'in'
    plt.rcParams['ytick.direction'] = 'in'
    '''------------------------------------------- Model validation ---------------------------------------------'''
    data_xy = np.load('/home/aistudio/work/NACA0010-65/level1/1_Fluent_results/level1_fluent_results.npy')[:, 0:2]
    test_data = data_xy

    # Model predictions generation
    u = model.predict(test_data)[:, 0] * a0
    v = model.predict(test_data)[:, 1] * a0
    p = model.predict(test_data)[:, 2] * p0
    rho = model.predict(test_data)[:, 3] * rho0

    airfoil_plot = np.load('/home/aistudio/work/NACA0010-65/level1/0_Mesh_files/ordered_airfoil_points.npy')

    # plot
    fig, ax = plt.subplots(figsize=(10, 4), dpi=300)
    ax.tricontourf(data_xy[:, 0], data_xy[:, 1], (u**2 + v**2)**0.5, levels=1000, cmap="coolwarm", vmin=48.9, vmax=462.5)
    plt.axis('scaled')
    plt.fill(airfoil_plot[:, 0], airfoil_plot[:, 1], color = 'w')
    norm =mpl.colors.Normalize(vmin=48.9, vmax=462.5)
    plt.rcParams['ytick.direction'] = 'out'
    plt.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap="coolwarm"), ax=ax, label='[m/s]')
    plt.rcParams['ytick.direction'] = 'in'    
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title('Velocity Magnitude')
    plt.show()   
    
    fig, ax = plt.subplots(figsize=(10, 4), dpi=300)
    ax.tricontourf(data_xy[:, 0], data_xy[:, 1], p, levels=1000, cmap="coolwarm", vmin=19587.33, vmax=98563.27)
    plt.axis('scaled')
    plt.fill(airfoil_plot[:, 0], airfoil_plot[:, 1], color = 'w')
    norm =mpl.colors.Normalize(vmin=19587.33, vmax=98563.27)
    plt.rcParams['ytick.direction'] = 'out'
    plt.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap="coolwarm"), ax=ax, label='[Pa]')
    plt.rcParams['ytick.direction'] = 'in'    
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title('Pressure')
    plt.show()  
    
    fig, ax = plt.subplots(figsize=(10, 4), dpi=300)
    ax.tricontourf(data_xy[:, 0], data_xy[:, 1], rho, levels=1000, cmap="coolwarm", vmin=0.315, vmax=1.109)
    plt.axis('scaled')
    plt.fill(airfoil_plot[:, 0], airfoil_plot[:, 1], color = 'w')
    norm =mpl.colors.Normalize(vmin=0.315, vmax=1.109)
    plt.rcParams['ytick.direction'] = 'out'
    plt.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap="coolwarm"), ax=ax, label='[kg/m^3]')
    plt.rcParams['ytick.direction'] = 'in'    
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title('Density')
    plt.show()  

def plot_naca633418_flow_field(model):
    """------------------------------ Fluid parameters  -----------------------"""
    # 空气绝热指数
    k = 1.4
    # 用于无量纲化的速度、压力和密度，取入口处的声速值、压力值和密度值
    p0 = 73048 #(pa)
    rho0 = 0.9 #(kg/m^3)
    a0 = np.sqrt(k*p0/rho0) #(m/s)

    import matplotlib as mpl
    from matplotlib import rcParams
    config = {"font.size": 16}
    rcParams.update(config)
    plt.rcParams['xtick.direction'] = 'in'
    plt.rcParams['ytick.direction'] = 'in'
    '''------------------------------------------- Model validation ---------------------------------------------'''
    data_xy = np.load('/home/aistudio/work/NACA6334-18/level1/1_Fluent_results/level1_fluent_results.npy')[:, 0:2]
    test_data = data_xy

    # Model predictions generation
    u = model.predict(test_data)[:, 0] * a0
    v = model.predict(test_data)[:, 1] * a0
    p = model.predict(test_data)[:, 2] * p0
    rho = model.predict(test_data)[:, 3] * rho0

    airfoil_plot = np.load('/home/aistudio/work/NACA6334-18/level1/0_Mesh_files/ordered_airfoil_points.npy')

    # plot
    fig, ax = plt.subplots(figsize=(10, 4), dpi=300)
    ax.tricontourf(data_xy[:, 0], data_xy[:, 1], (u**2 + v**2)**0.5, levels=1000, cmap="coolwarm", vmin=36.5, vmax=467)
    plt.axis('scaled')
    plt.fill(airfoil_plot[:, 0], airfoil_plot[:, 1], color = 'w')
    norm =mpl.colors.Normalize(vmin=36.5, vmax=467)
    plt.rcParams['ytick.direction'] = 'out'
    plt.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap="coolwarm"), ax=ax, label='[m/s]')
    plt.rcParams['ytick.direction'] = 'in'    
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title('Velocity Magnitude')
    plt.show()
    
    fig, ax = plt.subplots(figsize=(10, 4), dpi=300)
    ax.tricontourf(data_xy[:, 0], data_xy[:, 1], p, levels=1000, cmap="coolwarm", vmin=19900, vmax=100000)
    plt.axis('scaled')
    plt.fill(airfoil_plot[:, 0], airfoil_plot[:, 1], color = 'w')
    norm =mpl.colors.Normalize(vmin=19900, vmax=100000)
    plt.rcParams['ytick.direction'] = 'out'
    plt.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap="coolwarm"), ax=ax, label='[Pa]')
    plt.rcParams['ytick.direction'] = 'in'    
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title('Pressure')
    plt.show()
    
    fig, ax = plt.subplots(figsize=(10, 4), dpi=300)
    ax.tricontourf(data_xy[:, 0], data_xy[:, 1], rho, levels=1000, cmap="coolwarm", vmin=0.341, vmax=1.13)
    plt.axis('scaled')
    plt.fill(airfoil_plot[:, 0], airfoil_plot[:, 1], color = 'w')
    norm =mpl.colors.Normalize(vmin=0.341, vmax=1.13)
    plt.rcParams['ytick.direction'] = 'out'
    plt.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap="coolwarm"), ax=ax, label='[kg/m^3]')
    plt.rcParams['ytick.direction'] = 'in'    
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title('Density')
    plt.show()
    